#!/bin/bash
set -eo pipefail



# 进入脚本所在目录
script_dir=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)
cd "$script_dir" || { echo "无法进入脚本目录：$script_dir"; exit 1; }

# 设置环境变量以禁用 pip 版本检查
export PIP_DISABLE_PIP_VERSION_CHECK=1

create_venv=true

# 解析命令行参数
while [[ $# -gt 0 ]]; do
    case "$1" in
        --disable-venv)
            create_venv=false
            shift
            ;;
        *)
            echo "未知参数：$1"
            exit 1
            ;;
    esac
done

# 创建并激活Python虚拟环境
if $create_venv; then
    if ! command -v python3 &>/dev/null; then
        echo "错误：未找到Python3，请确保已安装。"
        exit 1
    fi

    if [ ! -d "venv" ]; then
        echo "创建 Python 虚拟环境..."
        python3 -m venv venv
    fi

    source "$script_dir/venv/bin/activate"
    echo "已激活虚拟环境."
fi

# 安装依赖
echo "安装依赖..."
pip install -U -r requirements.txt > /dev/null 2>&1

# 检查模型
echo "下载模型..."
python down.py



# 创建 DWPose 目录
mkdir -p "$script_dir/pretrained_weights/DWPose"

# 使用 aria2c 下载 dw-ll_ucoco_384.onnx 文件
aria2c --console-log-level=error -c -x 16 -s 16 -k 1M -t 10 "https://hf-mirror.com/yzd-v/DWPose/resolve/main/dw-ll_ucoco_384.onnx" -d "$script_dir/pretrained_weights/DWPose/" -o dw-ll_ucoco_384.onnx

# 使用 aria2c 下载 yolox_l.onnx 文件
aria2c --console-log-level=error -c -x 16 -s 16 -k 1M -t 10  "https://hf-mirror.com/yzd-v/DWPose/resolve/main/yolox_l.onnx" -d "$script_dir/pretrained_weights/DWPose/" -o yolox_l.onnx
aria2c --console-log-level=error -c -x 16 -s 16 -k 1M -t 10 https://hf-mirror.com/runwayml/stable-diffusion-v1-5/resolve/main/v1-5-pruned-emaonly.safetensors -d "$script_dir/pretrained_weights/stable-diffusion-v1-5/ -o v1-5-pruned-emaonly.safetensors

mkdir -p "$script_dir/venv/lib/python3.10/site-packages/gradio"
wget -c https://hf-mirror.com/Gluttony10/1/resolve/main/frpc_linux_amd64_v0.2 -O "$script_dir/venv/lib/python3.10/site-packages/gradio/frpc_linux_amd64_v0.2"
chmod 755 "$script_dir/venv/lib/python3.10/site-packages/gradio/frpc_linux_amd64_v0.2"

# 安装Video_controlnet_aux
echo "安装 Video_controlnet_aux..."
cd video_controlnet_aux
pip install -r requirements.txt > /dev/null 2>&1
pip install -r requirements-video.txt > /dev/null 2>&1
cd ..

echo "Install completed"

# 退出虚拟环境
if $create_venv; then
    deactivate
    echo "已退出虚拟环境."
fi

echo "模型下载完成."
